library(haven)
library(pscl)
library(dplyr)
library(ggplot2)
library(R2jags)
crash<-read.csv("/tx_intersection_crash.csv")

## Standardize the selected variables ##
crash_std <- as.data.frame(scale(crash[c("sidewalk_lenght_150ft_ft","lanes_major","lanes_minor","median_major","median_minor","on_system","median_width_ft_major","median_width_ft_minor",
                                         "lane_width_ft_major","lane_width_ft_minor","shoulder_width_ft_major","shoulder_width_ft_minor","aadt_lane_major","aadt_lane_minor",
                                         "truck_perc_major","truck_perc_minor","tot_WMT_sqmi","tot_WMT_pop","tot_WMT","speed_lim_mph_major","speed_lim_mph_minor",
                                         "f_local_major","f_local_minor","f_collector_major","f_collector_minor","f_arterial_major","f_arterial_minor",
                                         "a_rural","a_small_urban","a_urbanized","signalized_ind","approaches","dist_near_school_mi","dist_near_hops_mi","transit_ind","transit_stops_025mi_count",
                                         "pop_den","employ_den","cen_tr_income","cen_tr_age","county_rain")]))
crash_ind <- data.frame(crash[c("tot_crash_count")])
## Bind the columns of the two dataset ##
crash_std <- dplyr::bind_cols(crash_std, crash_ind)
## Remove missing values in the combined dataset ##
crash_std <- na.omit(crash_std)

## Divide the dataset into 70% training data and 30% test data ##
## If the model is running too long, try running the model on a random sample of the data: dplyr::sample_n(crash_std, 150000, replace=FALSE) ##
set.seed(1234)
library(caret)
indexes = createDataPartition(crash_std$tot_crash_count, p = .7, list = F)
train = crash_std[indexes, ]
test_x = crash_std[-indexes, -42]
test_y = crash_std[-indexes, 42]

# JAGS model
N<-nrow(train)
y<-train$tot_crash_count
x1<-train$sidewalk_lenght_150ft_ft
x2<-train$lanes_major
x3<-train$lanes_minor
x4<-train$median_major
x5<-train$median_minor
x6<-train$on_system
x7<-train$median_width_ft_major
x8<-train$median_width_ft_minor
x9<-train$lane_width_ft_major
x10<-train$lane_width_ft_minor
x11<-train$shoulder_width_ft_major
x12<-train$shoulder_width_ft_minor
x13<-train$aadt_lane_major
x14<-train$aadt_lane_minor
x15<-train$truck_perc_major
x16<-train$truck_perc_minor
x17<-train$tot_WMT_sqmi
x18<-train$tot_WMT_pop
x19<-train$tot_WMT
x20<-train$speed_lim_mph_major
x21<-train$speed_lim_mph_minor
x22<-train$f_local_major
x23<-train$f_local_minor
x24<-train$f_collector_major
x25<-train$f_collector_minor
x26<-train$f_arterial_major
x27<-train$f_arterial_minor
x28<-train$a_rural
x29<-train$a_small_urban
x30<-train$a_urbanized
x31<-train$signalized_ind
x32<-train$approaches
x33<-train$dist_near_school_mi
x34<-train$dist_near_hops_mi
x35<-train$transit_ind
x36<-train$transit_stops_025mi_count
x37<-train$pop_den
x38<-train$employ_den
x39<-train$cen_tr_income
x40<-train$cen_tr_age
x41<-train$county_rain

crash_tx.dat<-list("N","y","x1","x2","x3","x4","x5","x6","x7","x8","x9","x10","x11","x12","x13","x14","x15","x16","x17","x18","x19","x20","x21","x22","x23","x24","x25","x26","x27","x28","x29","x30","x31","x32","x33","x34","x35","x36","x37","x38","x39","x40","x41")

crash_tx <- function()
{
  for(i in 1:N)
  {
    # Likelihood
    y[i] ~ dnegbin(prob[i], phi)
    prob[i] <- phi/(phi+eps[i]*mu[i])
    log(mu[i]) <- b1*x1[i]+b2*x2[i]+b3*x3[i]+b4*x4[i]+b5*x5[i]+b6*x6[i]+b7*x7[i]+b8*x8[i]+b9*x9[i]+b10*x10[i]+b11*x11[i]+b12*x12[i]+b13*x13[i]+b14*x14[i]+b15*x15[i]+b16*x16[i]+b17*x17[i]+b18*x18[i]+b19*x19[i]+b20*x20[i]+b21*x21[i]+b22*x22[i]+b23*x23[i]+b24*x24[i]+b25*x25[i]+b26*x26[i]+b27*x27[i]+b28*x28[i]+b29*x29[i]+b30*x30[i]+b31*x31[i]+b32*x32[i]+b33*x33[i]+b34*x34[i]+b35*x35[i]+b36*x36[i]+b37*x37[i]+b38*x38[i]+b39*x39[i]+b40*x40[i]+b41*x41[i]   ## Dropping b0
    
    eps[i] ~ dgamma(f[i],t)
    f[i] <- 1+z[i] 
    z[i] ~ dbern(k)  
  }
  
  # Priors
  b1 ~ dnorm(0, 0.1)
  b2 ~ dnorm(0, 0.1)
  b3 ~ dnorm(0, 0.1)
  b4 ~ dnorm(0, 0.1)
  b5 ~ dnorm(0, 0.1)
  b6 ~ dnorm(0, 0.1)
  b7 ~ dnorm(0, 0.1)
  b8 ~ dnorm(0, 0.1)
  b9 ~ dnorm(0, 0.1)
  b10 ~ dnorm(0, 0.1)
  b11 ~ dnorm(0, 0.1)
  b12 ~ dnorm(0, 0.1)
  b13 ~ dnorm(0, 0.1)
  b14 ~ dnorm(0, 0.1)
  b15 ~ dnorm(0, 0.1)
  b16 ~ dnorm(0, 0.1)
  b17 ~ dnorm(0, 0.1)
  b18 ~ dnorm(0, 0.1)
  b19 ~ dnorm(0, 0.1)
  b20 ~ dnorm(0, 0.1)
  b21 ~ dnorm(0, 0.1)
  b22 ~ dnorm(0, 0.1)
  b23 ~ dnorm(0, 0.1)
  b24 ~ dnorm(0, 0.1)
  b25 ~ dnorm(0, 0.1)
  b26 ~ dnorm(0, 0.1)
  b27 ~ dnorm(0, 0.1)
  b28 ~ dnorm(0, 0.1)
  b29 ~ dnorm(0, 0.1)
  b30 ~ dnorm(0, 0.1)
  b31 ~ dnorm(0, 0.1)
  b32 ~ dnorm(0, 0.1)
  b33 ~ dnorm(0, 0.1)
  b34 ~ dnorm(0, 0.1)
  b35 ~ dnorm(0, 0.1)
  b36 ~ dnorm(0, 0.1)
  b37 ~ dnorm(0, 0.1)
  b38 ~ dnorm(0, 0.1)
  b39 ~ dnorm(0, 0.1)
  b40 ~ dnorm(0, 0.1)
  b41 ~ dnorm(0, 0.1)
  
  k ~ dbeta(1, 1)
  t <- (1-k)/k
  mean.eps <- (t+2)/(t*(t+1))
  log.mean.eps <- log(mean.eps) ## Calculating b0
  
  phi ~ dgamma(0.1, 0.1)
}


mon<-c("b1","b2","b3","b4","b5","b6","b7","b8","b9","b10","b11","b12","b13","b14","b15","b16","b17","b18","b19","b20","b21","b22","b23","b24","b25","b26","b27","b28","b29","b30","b31","b32","b33","b34","b35","b36","b37","b38","b39","b40","b41","log.mean.eps")
crash_tx.out <- jags(data=crash_tx.dat, inits=NULL, parameters.to.save=mon, model.file=crash_tx, n.iter = 1000, n.burn=100, n.thin=10)
update(crash_tx.out, 1000)

# Predictions
crash_tx.mat <- as.mcmc(crash_tx.out)
crash_tx.mat <- as.matrix(crash_tx.mat)
beta <- as.matrix(apply(crash_tx.mat[,1:41],2,mean))
X <- as.matrix(test_x[,c(1,10,11,12,13,14,15,16,17,18,19,2,20,21,22,23,24,25,26,27,28,29,3,30,31,32,33,34,35,36,37,38,39,4,40,41,5,6,7,8,9)])
pred_y <- exp(X%*%beta)
pred_y<-as.mcmc(pred_y)
plot(density(pred_y))

## filter pred_y < 2000 ##
## Compute R-squared ##
res <- test_y - pred_y
1 - var(res) / var(test_y)

## Compute the predictions of the model, and calculate MSE, MAE, and RMSE ##
mse = mean((test_y - pred_y)^2)
mae = caret::MAE(test_y, pred_y)
rmse = caret::RMSE(test_y, pred_y)
cat("MSE: ", mse, "MAE: ", mae, " RMSE: ", rmse)


## Since the dataset is imbalanced (About 70% of the intersections have 0 crashes), I try to balance the dataset by undersampling intersections with 0 crashes, and oversampling intersections with non-zero crashes ##
## The balanced dataset should have about 50% zero-crash intersections and 50% non-zero-crash intersections ##
## Data over and undersampling ##
crash_data <- as.data.frame(crash[c("sidewalk_lenght_150ft_ft","lanes_major","lanes_minor","median_major","median_minor","on_system","median_width_ft_major","median_width_ft_minor",
                                         "lane_width_ft_major","lane_width_ft_minor","shoulder_width_ft_major","shoulder_width_ft_minor","aadt_lane_major","aadt_lane_minor",
                                         "truck_perc_major","truck_perc_minor","tot_WMT_sqmi","tot_WMT_pop","tot_WMT","speed_lim_mph_major","speed_lim_mph_minor",
                                         "f_local_major","f_local_minor","f_collector_major","f_collector_minor","f_arterial_major","f_arterial_minor",
                                         "a_rural","a_small_urban","a_urbanized","signalized_ind","approaches","dist_near_school_mi","dist_near_hops_mi","transit_ind","transit_stops_025mi_count",
                                         "pop_den","employ_den","cen_tr_income","cen_tr_age","county_rain","tot_crash_count")])
## Remove missing values in the combined dataset ##
crash_data <- na.omit(crash_data)

library(ROSE)
crash_data$label <- NA
crash_data$label[which(crash_data$tot_crash_count==0)] <- 0
crash_data$label[which(crash_data$tot_crash_count!=0)] <- 1
crash_balance <- ovun.sample(label ~ ., data = crash_data, method = "both", seed = 1234)$data
table(crash_balance$label)
crash_balance <- crash_balance[,1:42]
crash_balance <- na.omit(crash_balance)

## Divide the dataset into 70% training data and 30% test data ##
## If the model is running too long, try running the model on a random sample of the data: dplyr::sample_n(crash_balance, 150000, replace=FALSE) ##
set.seed(1234)
library(caret)
indexes = createDataPartition(crash_std$tot_crash_count, p = .7, list = F)
train = crash_std[indexes, ]
test_x = crash_std[-indexes, -42]
test_y = crash_std[-indexes, 42]

# JAGS model
N<-nrow(train)
y<-train$tot_crash_count
x1<-train$sidewalk_lenght_150ft_ft
x2<-train$lanes_major
x3<-train$lanes_minor
x4<-train$median_major
x5<-train$median_minor
x6<-train$on_system
x7<-train$median_width_ft_major
x8<-train$median_width_ft_minor
x9<-train$lane_width_ft_major
x10<-train$lane_width_ft_minor
x11<-train$shoulder_width_ft_major
x12<-train$shoulder_width_ft_minor
x13<-train$aadt_lane_major
x14<-train$aadt_lane_minor
x15<-train$truck_perc_major
x16<-train$truck_perc_minor
x17<-train$tot_WMT_sqmi
x18<-train$tot_WMT_pop
x19<-train$tot_WMT
x20<-train$speed_lim_mph_major
x21<-train$speed_lim_mph_minor
x22<-train$f_local_major
x23<-train$f_local_minor
x24<-train$f_collector_major
x25<-train$f_collector_minor
x26<-train$f_arterial_major
x27<-train$f_arterial_minor
x28<-train$a_rural
x29<-train$a_small_urban
x30<-train$a_urbanized
x31<-train$signalized_ind
x32<-train$approaches
x33<-train$dist_near_school_mi
x34<-train$dist_near_hops_mi
x35<-train$transit_ind
x36<-train$transit_stops_025mi_count
x37<-train$pop_den
x38<-train$employ_den
x39<-train$cen_tr_income
x40<-train$cen_tr_age
x41<-train$county_rain

crash_tx.dat<-list("N","y","x1","x2","x3","x4","x5","x6","x7","x8","x9","x10","x11","x12","x13","x14","x15","x16","x17","x18","x19","x20","x21","x22","x23","x24","x25","x26","x27","x28","x29","x30","x31","x32","x33","x34","x35","x36","x37","x38","x39","x40","x41")

crash_tx <- function()
{
  for(i in 1:N)
  {
    # Likelihood
    y[i] ~ dnegbin(prob[i], phi)
    prob[i] <- phi/(phi+eps[i]*mu[i])
    log(mu[i]) <- b1*x1[i]+b2*x2[i]+b3*x3[i]+b4*x4[i]+b5*x5[i]+b6*x6[i]+b7*x7[i]+b8*x8[i]+b9*x9[i]+b10*x10[i]+b11*x11[i]+b12*x12[i]+b13*x13[i]+b14*x14[i]+b15*x15[i]+b16*x16[i]+b17*x17[i]+b18*x18[i]+b19*x19[i]+b20*x20[i]+b21*x21[i]+b22*x22[i]+b23*x23[i]+b24*x24[i]+b25*x25[i]+b26*x26[i]+b27*x27[i]+b28*x28[i]+b29*x29[i]+b30*x30[i]+b31*x31[i]+b32*x32[i]+b33*x33[i]+b34*x34[i]+b35*x35[i]+b36*x36[i]+b37*x37[i]+b38*x38[i]+b39*x39[i]+b40*x40[i]+b41*x41[i]   ## Dropping b0
    
    eps[i] ~ dgamma(f[i],t)
    f[i] <- 1+z[i] 
    z[i] ~ dbern(k)  
  }
  
  # Priors
  b1 ~ dnorm(0, 0.1)
  b2 ~ dnorm(0, 0.1)
  b3 ~ dnorm(0, 0.1)
  b4 ~ dnorm(0, 0.1)
  b5 ~ dnorm(0, 0.1)
  b6 ~ dnorm(0, 0.1)
  b7 ~ dnorm(0, 0.1)
  b8 ~ dnorm(0, 0.1)
  b9 ~ dnorm(0, 0.1)
  b10 ~ dnorm(0, 0.1)
  b11 ~ dnorm(0, 0.1)
  b12 ~ dnorm(0, 0.1)
  b13 ~ dnorm(0, 0.1)
  b14 ~ dnorm(0, 0.1)
  b15 ~ dnorm(0, 0.1)
  b16 ~ dnorm(0, 0.1)
  b17 ~ dnorm(0, 0.1)
  b18 ~ dnorm(0, 0.1)
  b19 ~ dnorm(0, 0.1)
  b20 ~ dnorm(0, 0.1)
  b21 ~ dnorm(0, 0.1)
  b22 ~ dnorm(0, 0.1)
  b23 ~ dnorm(0, 0.1)
  b24 ~ dnorm(0, 0.1)
  b25 ~ dnorm(0, 0.1)
  b26 ~ dnorm(0, 0.1)
  b27 ~ dnorm(0, 0.1)
  b28 ~ dnorm(0, 0.1)
  b29 ~ dnorm(0, 0.1)
  b30 ~ dnorm(0, 0.1)
  b31 ~ dnorm(0, 0.1)
  b32 ~ dnorm(0, 0.1)
  b33 ~ dnorm(0, 0.1)
  b34 ~ dnorm(0, 0.1)
  b35 ~ dnorm(0, 0.1)
  b36 ~ dnorm(0, 0.1)
  b37 ~ dnorm(0, 0.1)
  b38 ~ dnorm(0, 0.1)
  b39 ~ dnorm(0, 0.1)
  b40 ~ dnorm(0, 0.1)
  b41 ~ dnorm(0, 0.1)
  
  k ~ dbeta(1, 1)
  t <- (1-k)/k
  mean.eps <- (t+2)/(t*(t+1))
  log.mean.eps <- log(mean.eps) ## Calculating b0
  
  phi ~ dgamma(0.1, 0.1)
}


mon<-c("b1","b2","b3","b4","b5","b6","b7","b8","b9","b10","b11","b12","b13","b14","b15","b16","b17","b18","b19","b20","b21","b22","b23","b24","b25","b26","b27","b28","b29","b30","b31","b32","b33","b34","b35","b36","b37","b38","b39","b40","b41","log.mean.eps")
crash_tx.out <- jags(data=crash_tx.dat, inits=NULL, parameters.to.save=mon, model.file=crash_tx, n.iter = 1000, n.burn=100, n.thin=10)
update(crash_tx.out, 500)

crashtx<-as.mcmc(crash_tx.out)
geweke.diag(crashtx, frac1=0.1, frac2=0.5)
geweke.diag(crashtx, frac1=0.2, frac2=0.4)
heidel.diag(crashtx)
traceplot(crashtx)

# Predictions
crash_tx.mat <- as.mcmc(crash_tx.out)
crash_tx.mat <- as.matrix(crash_tx.mat)
beta <- as.matrix(apply(crash_tx.mat[,1:41],2,mean))
X <- as.matrix(test_x[,c(1,10,11,12,13,14,15,16,17,18,19,2,20,21,22,23,24,25,26,27,28,29,3,30,31,32,33,34,35,36,37,38,39,4,40,41,5,6,7,8,9)])
pred_y <- exp(X%*%beta)
pred_y<-as.mcmc(pred_y)
plot(density(pred_y))

## filter pred_y < 2000 ##
## Compute R-squared ##
res <- test_y - pred_y
1 - var(res) / var(test_y)

## Compute the predictions of the model, and calculate MSE, MAE, and RMSE ##
mse = mean((test_y - pred_y)^2)
mae = caret::MAE(test_y, pred_y)
rmse = caret::RMSE(test_y, pred_y)
cat("MSE: ", mse, "MAE: ", mae, " RMSE: ", rmse)